
import logging

from .shellcode_factory import ShellcodeFactory
from .exceptions import CannotExploit

l = logging.getLogger("rex.exploit.exploit_factory")


class ExploitFactory:
    """
    Managing and creating exploits.
    """

    def __init__(self, crash, blacklist_techniques=None, **kwargs):
        """
        :param crash: an exploitable crash object
        :param shellcode_opts: options to pass to the ShellcodeFactory
        :param blacklist_techniques: a set of techniques to skip
        """

        if not crash.exploitable():
            raise CannotExploit("crash cannot be exploited")

        self.crash = crash
        self.binary = crash.binary
        self.os = crash.project.loader.main_object.os

        # let's put together our rop gadgets
        self.rop = self.crash.rop

        # and let's gather some shellcode
        shellcode_opts = kwargs.pop('shellcode_opts', {})
        self.shellcode = ShellcodeFactory(crash.project, **shellcode_opts)

        self.arsenal = { }

        self.blacklist_techniques = set() if blacklist_techniques is None else blacklist_techniques
        self.tech_opts = kwargs

    def dump(self):
        for exploit_name in self.arsenal:
            self.arsenal[exploit_name].script(filename="%s.py" % exploit_name)
        return len(self.arsenal)

    def initialize(self):

        os = 'unix' if self.os.startswith('UNIX') else self.os

        for technique in Techniques[os]:
            p = technique(self.crash, self.rop, self.shellcode)
            if p.name in self.blacklist_techniques:
                l.info("[-] Skipping technique %s as it is in the blacklist.", p.name)
                continue
            if not p.check():
                l.info("[-] Skipping technique %s as the check fails.", p.name)
                continue
            try:
                l.info("[+] Attempt to apply technique %s...", p.name)
                self.arsenal[p.name] = p.apply(**self.tech_opts)
            except CannotExploit as ex:
                l.info("[-] Cannot apply technique %s: %s.", p.name, ex)


from .techniques import Techniques
